#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/8/12 22:02
@Author   : liujingmao
@File     : 3.BaseTool子类.py
"""
from typing import Any, Type

from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import BaseTool


class MultiplyInput(BaseModel):
    a: int = Field(..., description='第一个数字')
    b: int = Field(..., description='第二个数字')


class MultiplyTool(BaseTool):
    """two inters to multiply two numbers"""
    name = 'multiply_tool'
    description = "Two inters to multiply two numbers"
    args_schema: Type[BaseModel] = MultiplyInput

    def _run(self, *args: Any, **kwargs: Any) -> Any:
        """将传入的a和b相乘后返回"""
        return kwargs.get('a') * kwargs.get('b')


calculator = MultiplyTool()
print("name:", calculator.name)
print("description:", calculator.description)
print("args:", calculator.args)
print("是否返回:", calculator.return_direct)

print(calculator.invoke({"a": 2, "b": 3}))
